import numpy as np
from sklearn.linear_model import LinearRegression
from sklearn.model_selection import train_test_split

from typing import Tuple, Any

def train_model(X: np.ndarray, y: np.ndarray, test_size: float = 0.1, random_state: int = 114514) -> Tuple[Any, np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    '''训练线性回归模型
    
    Args:
        X: 特征矩阵
        y: 目标变量
        test_size: 测试集比例
        random_state: 随机种子
    
    Returns:
        model: 训练好的模型
        X_train: 训练集特征
        X_test: 测试集特征
        y_train: 训练集标签
        y_test: 测试集标签
    '''
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=test_size, random_state=random_state
    )
    
    model = LinearRegression()
    model.fit(X_train, y_train)
    
    return model, X_train, X_test, y_train, y_test

def predict(model: Any, X: np.ndarray) -> np.ndarray:
    '''预测
    
    Args:
        model: 训练好的模型
        X: 待预测的特征数据
    
    Returns:
        预测结果
    '''
    return model.predict(X) 